from transformers import T5ForConditionalGeneration, T5Tokenizer

class Translator:
    def __init__(self, model_name="/data/googlemadlad400-3b-mt/"):
        self.model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
        self.tokenizer = T5Tokenizer.from_pretrained(model_name)

    def translate(self, text, target_language_code):
        input_text = f"<{target_language_code}> {text}"
        input_ids = self.tokenizer(input_text, return_tensors="pt").input_ids.to(self.model.device)
        outputs = self.model.generate(input_ids=input_ids)
        translation = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return translation

# 示例用法
translator = Translator()
translated_text = translator.translate("我爱你!", "2en")
print(translated_text)  # 应该输出: Eu adoro pizza!







